import pandas as pd
import numpy as np


def mark_high(days_df, w_df):
    days_df = days_df.copy()
    w_df = w_df.copy()

    days_df['date'] = pd.to_datetime(days_df['date'])
    w_df['date'] = pd.to_datetime(w_df['date'])

    for df in [days_df, w_df]:
        iso = df['date'].dt.isocalendar()
        df['year_week'] = iso.year.astype(str) + '-' + iso.week.astype(str).str.zfill(2)

    days_df['dynamic_week_high'] = days_df.groupby(['code', 'year_week'])['high'].cummax()

    historical_high = w_df[['code', 'year_week', 'high']].rename(columns={'high': 'week_high'})
    dynamic_high = (
        days_df.sort_values('date')
        .groupby(['code', 'year_week'])
        .last()[['dynamic_week_high']]
        .reset_index()
        .rename(columns={'dynamic_week_high': 'week_high'})
    )

    combined_high = (
        pd.concat([historical_high, dynamic_high])
        .sort_values(['code', 'year_week'])
        .drop_duplicates(subset=['code', 'year_week'], keep='last')
    )

    windows = [10, 20, 30, 40]
    for window in windows:
        combined_high[f'high_{window}w'] = combined_high.groupby('code')['week_high'].transform(
            lambda s: s.shift(1).rolling(window, min_periods=1).max()
        )

    merged_df = pd.merge(
        days_df,
        combined_high[['code', 'year_week'] + [f'high_{window}w' for window in windows]],
        on=['code', 'year_week'],
        how='left'
    )
    # 生成特征 h_w10、h_w20、h_w30、h_w40
    for window in windows:
        merged_df[f'h_w{window}'] = ((merged_df['close']-merged_df[f'high_{window}w'] ) / merged_df[f'high_{window}w']).replace(0,1)*100

    drop_columns = ['year_week', 'dynamic_week_high'] + [f'high_{window}w' for window in windows]
    merged_df = merged_df.drop(columns=drop_columns)

    merged_df = cal_high_pct(merged_df)

    cols_to_drop = ['open', 'close', 'high', 'low', 'quote_rate', 
                   'high_limit', 'turnover', 't_rate','xl', 'xl_rate','gn_zt_ratio','hy_zt_ratio']
    merged_df = drop_base_columns(merged_df)
    return merged_df.round(2)


def cal_high_pct(df):
    df = df.copy().sort_values(by=['code', 'date'])
    
    # 生成特征 h_d20、h_d40、h_d60
    for window in [20, 40, 60]:
        high_col = f'high_{window}'
        df[high_col] = df.groupby('code')['high'].transform(
            lambda x: x.rolling(window=window, min_periods=1).max()
        )
        df[f'h_d{window}'] = (
            (df['close'] - df[high_col]) / df[high_col].replace(0, 1) * 100
        ).round(2)
    
    return df.drop(columns=[f'high_{w}' for w in [20,40,60]])

